#!/bin/bash

set -u

JETPACK_VER=
L4T_VER=
L4T_VER_MAJOR=
JETSON_MODULE=
JETSON_BOARD=
JETSON_CONF=
CMDARGS_ARRAY=

CMD_BSP="bsp"
CMD_SRC="src"
CMD_ROOTFS="rootfs"
CMD_FLASH="flash"

L4T_TOP=
CUR_TOP=$(dirname $(realpath $0))

################################################################

function usage()
{
	echo -e "\nJetson BSP Development tool"
	echo -e "Usage: $0 <OPTS> <CMD> [ARGS]"
	echo -e ""
	echo -e "  OPTS can be:"
	echo -e "    -h ------------- Show this usage, then exit"
	echo -e "    -l ------------- List all supported BSP Verions, Jetson Modules, Boards, then exit"
	echo -e "    -v <BSPVer> ---- JetPack or L4T BSP Version (essential)"
	echo -e "    -m <Module> ---- Jetson Module Name or Part Number (essential)"
	echo -e "    -b <Board> ----- Jetson Carrier Board Name or Number (essential)"
	echo -e "    -c <Conf> ------ Jetson Config File (optional)\n"
	echo -e "  CMD and ARGS can be:"
	echo -e "    bsp <args> ------- L4T BSP function"
	echo -e "       init ------------- Init the L4T BSP directory with specific string suffix"
	echo -e "       st --------------- Get current BSP status"
	echo -e "       release ---------- Release L4T BSP into a package"
	echo -e "       backup [args] ------- Backup Device images into a backup package"
	echo -e "       restore <pkg> [arg] - Restore the backup package to Device"
	echo -e "          [args] ------------- From where to backup and restore, can be: emmc, nvme"
	echo -e "    src <args> ------- L4T SRC Code function"
	echo -e "       build [args] ----- Build L4T SRC Kernel and DTB BSP"
	echo -e "          all -------------- Build L4T Kernel DTBs,Image,Modules"
	echo -e "          dtbs ------------- Build L4T Kernel DTBs"
	echo -e "          Image ------------ Build L4T Kernel Image"
	echo -e "          modules ---------- Build L4T Kernel Modules"
	echo -e "       install ---------- Install L4T Kernel SRC files (Image,DTB,Modules) into BSP"
	echo -e "       uninstall -------- Uninstall L4T Kernel SRC files (Image,DTB,Modules) from BSP"
	echo -e "       reinstall -------- Uninstall and Install L4T Kernel SRC files again"
	echo -e "    rootfs [args] ---- L4T rootfs function, call the script function defined in conf with \"RK_ROOTFS_TOOL\" tag"
	echo -e "       install ---------- Install Ritrontek L4T rootfs stuff"
	echo -e "       uninstall -------- Uninstall Ritrontek L4T rootfs stuff"
	echo -e "    flash [<device> [args]] -- L4T Flash function, default to flash whole BSP to mmcblk0p1"
	echo -e "       <device> can be one of the below:"
	echo -e "          emmc ------------- Flash to Jetson emmc"
	echo -e "          nvme ------------- Flash to Jetson to setup an external nvme drive"
	echo -e "       [args] ----------- Flash options directly pass to flash.sh\n"
}

function get_option()
{
	local OPTSTR="v:m:b:c:lh"
	local OPTS=$(getopt -o ${OPTSTR} -n "$0" -- "$@")
	if [[ $? != 0 ]];
	then
		usage
		exit 1
	fi

	eval set -- ${OPTS}

	while true; do
		case "$1" in
		"-h")
			usage
			exit 0
			;;
		"-l")
			conf2_list_vers
			conf2_list_modules
			conf2_list_boards
			shift 2
			exit 0
			;;
		"-v")
			JETPACK_VER=$(conf2_find_jetpack_ver $2)
			if [[ -z ${JETPACK_VER} ]]; then
				echo -e "[ERROR] Unsupport BSP Version: $2\n"
				exit 1
			fi
			L4T_VER=$(conf2_get_l4t_ver ${JETPACK_VER})
			L4T_VER_MAJOR=$(echo ${L4T_VER} | awk -F . '{print $1}')
			shift 2
			;;
		"-m")
			local module_pnum=$(conf2_find_module_pnum $2)
			if [[ -z ${module_pnum} ]]; then
				echo -e "[ERROR] Unsupport Module: $2\n"
				exit 1
			fi
			JETSON_MODULE=${module_pnum}
			shift 2
			;;
		"-b")
			local board_pnum=$2
			local module_pnum_list=$(conf2_get_module_pnum_list_with_board ${board_pnum})
			if [[ -z ${module_pnum_list} ]]; then
				if [[ ${board_pnum} =~ ^rk[a-z0-9]{6,7}$ ]]; then
					board_pnum="$2-v01"
					module_pnum_list=$(conf2_get_module_pnum_list_with_board ${board_pnum})
				fi
				if [[ -z ${module_pnum_list} ]]; then
					echo -e "[ERROR] Unsupport Board: $2\n"
					exit 1
				fi
			fi
			JETSON_BOARD=${board_pnum}
			shift 2
			;;
		"-c")
			JETSON_CONF="$2"
			shift 2
			;;
		"--")
			shift
			CMDARGS_ARRAY=($@)
			break
			;;
		*)
			echo -e "[ERROR] Option error\n"
			exit 1
			;;
		esac
	done

	if [[ -z ${JETPACK_VER} ]] \
		|| [[ -z ${JETSON_MODULE} ]] \
		|| [[ -z ${JETSON_BOARD} ]]; then

		echo -e "[ERROR] Option error\n"
		exit 1
	fi

	if [[ -z ${JETSON_CONF} ]]; then
		if [[ ${JETSON_BOARD} =~ ^p[0-9]{4}-[0-9]{4} ]]; then
			JETSON_CONF=${JETSON_BOARD}+${JETSON_MODULE}
		else
			JETSON_CONF=ritrontek_${JETSON_MODULE}_${JETSON_BOARD}
		fi
	fi

	if [[ ${#CMDARGS_ARRAY[@]} -le 0 ]]; then
		echo -e "[ERROR] No CMD specified\n"
		exit 1
	fi
}

function init_host_env()
{
	# NVIDIA Official Host Init script:
	# -> ${L4T_TOP}/tools/l4t_flash_prerequisites.sh

	# Check needed package whether installed
	type dtc &> /dev/null
	local ret=$?
	if [[ ${ret} -ne 0 ]]; then
		echo -e "[INFO] Install needed packages:"
		sudo_cmd "apt update && apt install -y python python3 \
					qemu-user-static \
					device-tree-compiler \
					pkg-config \
					libxml2-utils \
					simg2img \
					network-manager \
					abootimg \
					sshpass \
					flex \
					bison \
					libssl-dev \
					libncurses-dev \
					nfs-kernel-server"
		ret=$?
		if [[ ${ret} -ne 0 ]]; then
			echo -e "[ERROR] Install packages failed\n"
			exit 1
		fi
	fi

	# Check dtc version >=1.5.0, upgrade it if needed
	local dtc_ver=$(dtc -v | awk '{print $3}' | sed 's/\.//g' 2>/dev/null)
	ret=$?
	if [[ ${ret} -ne 0 || ${dtc_ver} -lt 150 ]]; then
		echo -e "[WARN] \"dtc\" version is lower than 1.5.0, now install it\n"

		libftd_pkg="libfdt1_1.5.1-1_amd64.deb"
		dtc_pkg="device-tree-compiler_1.5.1-1_amd64.deb"

		if [[ ! -f ${libftd_pkg} ]]; then
			wget https://mirrors.aliyun.com/ubuntu/pool/main/d/device-tree-compiler/libfdt1_1.5.1-1_amd64.deb
		fi
		if [[ ! -f ${dtc_pkg} ]]; then
			wget https://mirrors.aliyun.com/ubuntu/pool/main/d/device-tree-compiler/device-tree-compiler_1.5.1-1_amd64.deb
		fi

		sudo_cmd "apt install -y ./${libftd_pkg} ./${dtc_pkg}"
		ret=$?
		if [[ ${ret} -ne 0 ]]; then
			echo -e "[ERROR] Download and install \"device-tree-compiler 1.5.1\" error:"
			echo -e ">>> ${libftd_pkg}"
			echo -e ">>> ${dtc_pkg}\n"
			exit 1
		fi
		rm ${libftd_pkg} ${dtc_pkg}

		dtc_ver=$(dtc -v | awk '{print $3}' | sed 's/\.//g' 2>/dev/null)
		ret=$?
		if [[ ${ret} -ne 0 || ${dtc_ver} -lt 150 ]]; then
			echo -e "[ERROR] \"dtc\" version cannot lower than 1.5.0.\n"
			exit 1
		fi
	fi

	# Check udisk2.service whether stopped
	local service_status=$(systemctl is-active udisks2.service)
	if [[ ${service_status} == "active" ]]; then
		sudo_cmd "systemctl stop udisks2.service"
	fi
}

################################################################

source ${CUR_TOP}/helper/helper.sh
source ${CUR_TOP}/helper/helper_conf.sh
source ${CUR_TOP}/helper/helper_sudo.sh
source ${CUR_TOP}/helper/helper_traverse.sh

init_host_env

get_option $@
L4T_TOP=${HOME}/nvidia/nvidia_sdk/JetPack${JETPACK_VER}_L4T${L4T_VER}_BSP/Linux_for_Tegra

echo -e "[DEBUG] CMDARGS_ARRAY=${CMDARGS_ARRAY[@]}"
CMD=${CMDARGS_ARRAY[0]}
ARGS=${CMDARGS_ARRAY[@]:1}
case ${CMD} in
	${CMD_BSP} )
		source ${CUR_TOP}/helper/helper_cmd_bsp.sh
		cmd_bsp ${ARGS}
		;;
	${CMD_SRC} )
		source ${CUR_TOP}/helper/helper_cmd_src.sh
		cmd_src ${ARGS}
		;;
	${CMD_ROOTFS} )
		source ${CUR_TOP}/helper/helper_cmd_rootfs.sh
		cmd_rootfs ${ARGS}
		;;
	${CMD_FLASH} )
		source ${CUR_TOP}/helper/helper_cmd_flash.sh
		cmd_flash ${ARGS}
		;;
	* )
		echo -e "[ERROR] Unsupport CMD: ${CMD}\n"
		exit 1
		;;
esac
